import torch
import torchvision.models as models
from torchvision.transforms import ToTensor
from PIL import Image
import time

#定义测试图像
print('加载测试图像。。。')
image = Image.new('RGB', (224, 224),color=(255,0,0))
transform = ToTensor()
input_image = transform(image).unsqueeze(0)
print(f'图像尺寸:{input_image.shape}')
#加载预训练模型
model_fp32 = models.resnet18(pretrained=True)
model_fp32.eval()
start = time.time()
with torch.no_grad():
    output_fp32 = model_fp32(input_image)

end = time.time()
print(f'Fp32模型推理时间为{end-start}')

# 动态量化
print('开始量化...')
model_dynamic = torch.quantization.quantize_dynamic(
    model_fp32,
    {torch.nn.Linear},
    dtype=torch.qint8,
)
print('动态量化完成。。。')
start = time.time()
with torch.no_grad():
    output_dynamic = model_dynamic(input_image)
end = time.time()
print(f'模型推理时间为{end-start}')

#静态量化
print('静态量化开始。。。')
model_static = models.resnet18(pretrained=True)
model_static.eval()
#配置量化参数
model_static.qconfig = torch.quantization.get_default_qconfig("fbgemm")
#准备量化
model_prepared = torch.quantization.prepare_qat(model_static, inplace=True)
model_quantized = torch.quantization.convert(model_prepared,inplace=True)
start = time.time()

with torch.no_grad():
    out = model_quantized(input_image)
end = time.time()
print(f'静态量化模型推理时间为{end-start}')
